import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm

def print_agent(agent, env, action_meaning, disaster=[], end=[]):    
    for y in range(env.nrow):
        for x in range(env.ncol):
            pos = get_env_pos(env, x, y)
            if pos in disaster:
                print('****', end=' ')
            elif pos in end:
                print('EEEE', end=' ')
            else:
                a = agent.best_action(pos)
                pi_str = ''
                for i in range(len(action_meaning)):
                    pi_str += action_meaning[i] if a[i] > 0 else 'o'
                print(pi_str, end=' ')
        print()

def get_env_pos(env, x, y):
    return y*env.ncol + x


class CliffWalkingEnv:
    def __init__(self, ncol = 12, nrow=4):
        self.ncol = ncol
        self.nrow = nrow
        self.x = 0
        self.y = self.nrow - 1
        # 转移矩阵P,每个位置都是一个state, P[state][action] = [(p, next_state, reward, done)]

    def step(self, action):
        # 原点在左上角
        # 4中动作，change[0]:上, change[1]:下, change[2]:左, change[3]:右。
        change = [[0, -1], [0, 1], [-1, 0], [1, 0]]
        self.x = min(self.ncol - 1, max(0, self.x + change[action][0]))
        self.y = min(self.nrow - 1, max(0, self.y + change[action][1]))
        reward = -1
        done = False
        # 下一个位置在悬崖或者终点
        if self.isDone(self.x , self.y): # 下一个位置在悬崖或者目标
            done = True
            if self.isCliff(self.x, self.y): # 下一个位置在悬崖
                reward = -100
        return self.getPos(self.x, self.y), reward, done

    def getPos(self, x, y):
        return y*self.ncol + x

    def isDone(self, x, y):
        # 当前在悬崖或者目标位置
        return self.isEnd(x, y) or self.isCliff(x, y)

    def isEnd(self, x, y): 
        # 是否目标位置
        return y == self.nrow - 1 and x == self.ncol - 1

    def isCliff(self, x, y):
        # 是否是悬崖
        return y == self.nrow - 1 and x > 0 and x < self.ncol - 1

    def reset(self):
        self.x = 0
        self.y = self.nrow - 1
        return self.getPos(self.x, self.y)

class Sarsa:
    def __init__(self, ncol, nrow, epsilon, alpha, gamma, n_action=4):
        self.Q = np.zeros([nrow * ncol, n_action])    
        self.n_action = n_action
        self.alpha = alpha
        self.gamma = gamma
        self.epsilon = epsilon

    def take_action(self, state):
        if np.random.random() < self.epsilon:
            action = np.random.randint(self.n_action)
        else:
            action = np.argmax(self.Q[state])
        return action

    def best_action(self, state):
        Q_max = np.max(self.Q[state])
        a = [0 for _ in range(self.n_action)]
        for i in range(self.n_action):
            if self.Q[state, i] == Q_max:
                a[i] = 1
        return a

    def update(self, s0, a0, r, s1, a1, _done):
        td_error = r + self.gamma * self.Q[s1, a1] - self.Q[s0, a0]
        self.Q[s0, a0] += self.alpha * td_error



def sarsa_episodes(agent, env, num_episodes, show=True):
    return_list = []
    for i in range(10):
        with tqdm(total=int(num_episodes/10), desc='Iteration %d' % i) as pbar:
            for i_episode in range(int(num_episodes/10)):
                episode_return = 0
                state = env.reset()
                action = agent.take_action(state)
                done = False
                while not done:
                    next_state, reward, done = env.step(action)
                    next_action = agent.take_action(next_state)
                    episode_return += reward
                    agent.update(state, action, reward, next_state, next_action, done)
                    state = next_state
                    action = next_action

                return_list.append(episode_return)
                if (i_episode+1) % 10 == 0:
                    pbar.set_postfix({
                        'episode':
                        '%d' % (num_episodes / 10 * i + i_episode + 1),
                        'return':
                        '%.3f' % np.mean(return_list[-10:])
                    })
                pbar.update(1)

    if show:
        episodes_list = list(range(len(return_list)))
        plt.plot(episodes_list, return_list)
        plt.xlabel('Episodes')
        plt.ylabel('Returns')
        plt.title('Sarsa on {}'.format('Cliff Walking'))
        plt.show()

if __name__ == "__main__":
    np.random.seed(0)
    ncol = 12
    nrow = 4
    alpha = 0.1
    epsilon = 0.1
    gamma = 0.9
    agent = Sarsa(ncol, nrow, epsilon, alpha, gamma)
    env = CliffWalkingEnv(ncol, nrow)
    num_episodes = 500
    # show_sarsa_reawrd(agent, env, num_episodes)
        

    action_meaning = ['^', 'v', '<', '>']
    print('Sarsa算法最终收敛得到的策略为：')
    print_agent(agent, env, action_meaning, list(range(37, 47)), [47])

    class nstep_sarsa:
        def __init__(self, n, ncol, nrow, epsilon, alpha, gamma, n_action=4):
            self.Q = np.zeros([nrow * ncol, n_action])
            self.n_action = n_action
            self.alpha = alpha
            self.gamma = gamma
            self.epsilon = epsilon
            self.n = n
            self.state_list = []
            self.action_list = []
            self.reward_list = []

        def take_action(self, state):
            if np.random.random() < self.epsilon:
                action = np.random.randint(self.n_action)
            else:
                action = np.argmax(self.Q[state])
            return action

        def best_action(self, state):
            Q_max = np.max(self.Q[state])
            a = [0 for _ in range(self.n_action)]
            for i in range(self.n_action):
                if self.Q[state, i] == Q_max:
                    a[i] = 1
            return a

        def update(self, s0, a0, r, s1, a1, done):
            self.state_list.append(s0)
            self.action_list.append(a0)
            self.reward_list.append(r)
            if len(self.state_list) == self.n:
                G = self.Q[s1, a1]
                for i in reversed(range(self.n)):
                    G = self.gamma * G + self.reward_list[i]
                    if done and i > 0: # 如果是终止状态，也更新堆栈里的状态，相当于退化成了了n-1,n-2,...,1步更新
                        s = self.state_list[i]
                        a = self.action_list[i]
                        self.Q[s, a] += self.alpha * (G - self.Q[s, a])
                s = self.state_list.pop(0)
                a = self.action_list.pop(0)
                self.reward_list.pop(0)
                # 更新Q(s,a)
                self.Q[s, a] += self.alpha * (G - self.Q[s, a])

            if done: # 如果是终止状态，清空堆栈
                self.state_list = []
                self.action_list = []
                self.reward_list = []

    np.random.seed(0)
    n_step = 5
    alpha = 0.1
    epsilon = 0.1
    gamma = 0.9
    agent = nstep_sarsa(n_step, ncol, nrow, epsilon, alpha, gamma)
    num_episodes = 500
    sarsa_episodes(agent, env, num_episodes, False)
    action_meaning = ['^', 'v', '<', '>']
    print('5步Sarsa算法最终收敛得到的策略为：')
    print_agent(agent, env, action_meaning, list(range(37, 47)), [47])
